package com.gitee.simpleplus.utils;

import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

public class ScanClassUtil {

    public static Set<Class<?>> scan(String... pkgs){
    	Set<Class<?>> classes = new HashSet<>();
    	Arrays.stream(pkgs).forEach(pk -> classes.addAll(scan0(pk)));
    	return classes;
    }
	
    public static Set<Class<?>> scan0(String packageName){

    	Set<Class<?>> classes = new HashSet<>();
        try {
            ClassLoader loader = Thread.currentThread().getContextClassLoader();
            Enumeration<URL> urls = loader.getResources(packageName.replace('.', '/'));
            while (urls.hasMoreElements()) {
                URI uri = urls.nextElement().toURI();
                switch (uri.getScheme().toLowerCase()) {
                case "jar":
                    scanFromJarProtocol(loader, classes, uri.getRawSchemeSpecificPart());
                    break;
                case "file":
                    scanFromFileProtocol(loader, classes, uri.getPath(), packageName);
                    break;
                default:
                    throw new URISyntaxException(uri.getScheme(), "unknown schema " + uri.getScheme());
                }

            }
        } catch (URISyntaxException | IOException e) {
            e.printStackTrace();
        }

        return classes;
    }

	private static void scanFromJarProtocol(ClassLoader loader, Set<Class<?>> classes, String fullPath) {
		final String jar = fullPath.substring(0, fullPath.lastIndexOf('!'));
		final String parent = fullPath.substring(fullPath.lastIndexOf('!') + 2);

		try (JarInputStream jarReader = new JarInputStream(new URL(jar).openStream())) {
			JarEntry e = null;
			while ((e = jarReader.getNextJarEntry()) != null) {
				String className = e.getName();
				if (!e.isDirectory() && className.startsWith(parent) && className.endsWith(".class")
						&& !className.contains("$")) {
					className = className.replace('/', '.');
					classes.add(loadClass(loader, className));
				}
				jarReader.closeEntry();
			}
		} catch (Exception e1) {
			throw new RuntimeException(e1);
		}

	}
	
	private static void scanFromFileProtocol(ClassLoader loader, Set<Class<?>> classes, String dir, String packageName){
        File directory = new File(dir);
        File[] files = directory.listFiles();
        if (directory.isDirectory() && files != null) {
            for (File classFile : files) {
                if (!classFile.isDirectory() && classFile.getName().endsWith(".class")) {
                    String className = String.format("%s.%s", packageName, classFile.getName());
                    classes.add(loadClass(loader, className));
                }
            }
        }
    }


    private static Class<?> loadClass(ClassLoader loader, String classPath) {
        classPath = classPath.substring(0, classPath.length() - 6);
        try {
			return loader.loadClass(classPath);
		} catch (ClassNotFoundException e) {
			throw new RuntimeException(e);
		}
    }

}
